import torch
import torch.nn.functional as F
import torch.nn as nn
def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
    """3x3 convolution with padding"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
                     padding=dilation, groups=groups, bias=False, dilation=dilation)

def conv_3(in_planes, out_planes, stride=1, groups=1, dilation=1):
    """3x3 convolution with padding"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
                     padding=dilation, groups=groups, bias=False, dilation=dilation)

def conv1x1(in_planes, out_planes, stride=1):
    """1x1 convolution"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)


class CombinedModel(nn.Module):
    def __init__(self,block_list,choose,planes):
        super(CombinedModel, self).__init__()
        self.block1 = block_list[0]
        # self.concat_conv1 = conv1x1(planes[0][choose[0]],planes[0][choose[1]])
        self.block2 = block_list[1]
        # self.concat_conv2 = conv1x1(planes[1][choose[1]],planes[1][choose[1]])
        # self.concat_conv21 = conv1x1(planes[1][choose[1]],planes[1][choose[2]])
        # self.concat_conv22 = conv1x1(planes[1][choose[2]],planes[1][choose[2]])
        ######
        self.block2_concat1 = conv1x1(planes[0][choose[1]][0],planes[0][choose[1]][1])
        # self.block2_concat2 = conv1x1(planes[0][choose[1]][1],planes[0][choose[1]][2])
        self.block3_concat3 = conv1x1(planes[1][choose[2]][0],planes[1][choose[2]][1])
        # self.block3_concat4 = conv1x1(planes[1][choose[2]][1],planes[1][choose[2]][2])
        #####
        # self.concat_fc = nn.Linear(planes[1][choose[1]],planes[1][choose[2]])
        self.block3 = block_list[2]
        # self.concat_conv3 = conv1x1(planes[2][choose[2]],planes[2][choose[3]])

        self.block4 = block_list[3]
        # self.choose = choose
    def forward(self, x):
        x = self.block1(x)
        # x = self.concat_conv1(x)
        x = self.block2(x)
        # x = self.concat_conv2(x)
        x = self.block2_concat1(x)
        # x = self.block2_concat2(x)
        result = {'mid': x}
        # print(x.size())
        x = self.block3_concat3(x)
        # x = self.block3_concat4(x)
        # print(x.size())
        # x = self.concat_conv21(x)
        # x = self.concat_conv22(x)
        # x = self.concat_fc(x)

        x = self.block3(x)
        # x = self.concat_conv3(x)

        # if self.choose[3] != 2:
        #     for i in range(len(self.block4)):
        #         if i !=len(self.block4)-1:
        #             x = self.block4[i](x)
        #         else:
        #             # print(x.size())
        #             x = torch.flatten(x,1)

        #             x = self.block4[i](x)
        # else:
        #     x = self.block4(x)
        #     x = torch.flatten(x,1)
        x = self.block4(x)
        result['output'] = x
        return result

class CombinedSingleModel(nn.Module):
    def __init__(self,block_list,index):
        super(CombinedSingleModel, self).__init__()
        self.block1 = block_list[0]
        self.block2 = block_list[1]
        self.block3 = block_list[2]
        self.block4 = block_list[3]
        # self.choose = index
    def forward(self, x):
        x = self.block1(x)
        # print(self.block1)
        x = self.block2(x)
        # print(self.block2)
        result = {'mid':x}
        x = self.block3(x)
        x = self.block4(x)
        # print(x.size())
        # if self.choose != 2:
        #     for i in range(len(self.block4)):
        #         if i !=len(self.block4)-1:
        #             x = self.block4[i](x)
        #         else:
        #             x = torch.flatten(x,1)
        #             x = self.block4[i](x)
        # else:
        #     x = self.block4(x)
        #     x = torch.flatten(x,1)
        #     # print(x.size())
        result['output'] = x
        return result


class CombinedOneBlock(nn.Module):
    def __init__(self,block1,block2):
        super(CombinedOneBlock, self).__init__()
        self.block1 = block1
        self.block2 = block2
        # self.choose = index
    def forward(self, x):
        result = {'input':x}
        x = self.block1(x)
        x = self.block2(x)
        result['output'] = x
        return result